import json
from typing import Set, List, Iterable


class JsonLReader:

    def __init__(self, file: str, encoding="utf-8", keys: Set[str] = None, drop_list: Set[int] = None):
        """
        :param file: json line file path.
        :param encoding: default "utf-8"
        """
        self.file = None
        self.file = open(file, "r", encoding=encoding)
        self.keys = keys
        self.drop_list = drop_list

    def yield_read(self):
        has_drop = self.drop_list is not None
        index = -1
        while True:
            line = self.file.readline()
            index = index + 1
            if line is None or len(line) == 0:
                break
            else:
                if has_drop and index in self.drop_list:
                    continue
                yield self.json_load(line)

    def json_load(self, line) -> dict:
        if self.keys is None:
            data = json.loads(line)
        else:
            data = json.loads(line)
            for key in data:
                if key in data.keys() and key not in self.keys:
                    data.pop(key)
        return data

    def read(self) -> List[dict]:
        return [item for item in self.yield_read()]

    def read_idx(self, idx: int) -> dict:
        index = -1
        while index <= idx:
            line = self.file.readline()
            index = index + 1
            if index == idx:
                return self.json_load(line)

    def __del__(self):
        if self.file:
            self.file.close()


class JsonLWriter:
    def __init__(self, file: str, mode: str = "w+", encoding="utf-8"):
        self.file = None
        self.file = open(file, mode, encoding=encoding)

    def close(self):
        if self.file:
            self.file.close()

    def save_one(self, item: dict):
        self.file.write(self.dumps_line(item))

    def save(self, data: Iterable[dict]):
        data = [JsonLWriter.dumps_line(each) for each in data]
        self.file.writelines(data)

    @staticmethod
    def dumps_line(item: dict):
        return json.dumps(item) + "\n"


if __name__ == '__main__':
    pass
    from config.Config import BASE_DIR
    from os import path

    f = path.join(BASE_DIR, 'resource/data/context.jsonl')
    g = path.join(BASE_DIR, 'resource/data/context2.jsonl')
    json_reader = JsonLReader(f)
    json_write = JsonLWriter(g)
    d = json_reader.yield_read()
    json_write.save(d)
